package com.cnbot.aimolib.recognize;

import android.graphics.RectF;
import android.os.SystemClock;

import com.aimall.core.define.ImoImageFormat;
import com.aimall.core.define.ImoImageOrientation;
import com.cnbot.aimolib.AimoEngine;
import com.cnbot.aimolib.DetectionTrackerCore;
import com.cnbot.aimolib.bean.AimoFrameData;
import com.cnbot.aimolib.bean.AimoOutlineData;
import com.cnbot.aimolib.bean.AimoRawFaceData;
import com.cnbot.aimolib.db.AimoFaceDbData;
import com.cnbot.aimolib.db.AimoFaceDbManager;
import com.white.loglib.L;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * @author ytf
 * Created by on 2021/06/07 15:44
 */
public class RecognizeTask implements Runnable {

    private ArrayBlockingQueue<AimoFrameData> queue;
    private volatile AtomicBoolean isExit = new AtomicBoolean(false);
    private volatile boolean isRunning = false;

    private IRecognizeStrategy strategy;
    private Callback mCallback;
    private ImoImageFormat format;
    private ImoImageOrientation orientation;
    private Object[] transExtra;
    private float threshold;
    private long timeout;

    public RecognizeTask(Builder builder) {
        queue = new ArrayBlockingQueue<>(1);
        strategy = builder.strategy == null ? new DefaultRecognizeStrategy() : builder.strategy;
        mCallback = builder.mCallback;
        format = builder.format == null ? ImoImageFormat.IMO_IMAGE_NV21 : builder.format;
        orientation = builder.orientation == null ? ImoImageOrientation.IMO_IMAGE_UP : builder.orientation;
        transExtra = builder.transExtra;
        threshold = builder.faceThreshold == 0 ? 0.65f : builder.faceThreshold;
        timeout = (long) (builder.timeoutSeconds * 1000);
    }

    public boolean isActive(){
        return isRunning && !isExit.get();
    }

    public void post(AimoFrameData data){
        if(!isRunning) return;

        try {
            boolean ret = queue.offer(data);
            if(!ret){
//                L.e("队列已满，丢弃");
            }
        }catch (Exception e){

        }
    }

    @Override
    public void run() {
        DetectionTrackerCore core = AimoEngine.getInstance().getCore();
        if(!core.canWork()) {
            mCallback.onError("aimo not prepared");
            return;
        }

        List<AimoFaceDbData> allUsers = AimoFaceDbManager.getInstance().listAll();
        if(allUsers == null || allUsers.size() == 0){
            mCallback.onError("未识别，无人脸记录");
            return;
        }

        long start = SystemClock.uptimeMillis();

        RectF preview = new RectF(0f, 0f, 0f, 0f);
        //边缘触发控制
        boolean hasDetectComplete = false;

        while (!isExit.get()){
            isRunning = true;
            try {
                //2秒无数据自动退出
                AimoFrameData data = queue.poll(2, TimeUnit.SECONDS);
                if(data == null){
                    //超时没有数据，该结束了
                    exit();
                    mCallback.onError("无数据");
                }else if (data.isValid()) {
                    preview.right = data.width;
                    preview.bottom = data.height;

                    List<AimoRawFaceData> ret = core.detectFromStreamBytes(data.data, data.width, data.height, format, orientation);
                    if(ret != null && ret.size() > 0) {

                        List<AimoOutlineData> outlines = mapPoints(ret);
                        AimoFaceDbData exist = strategy.matchFace(ret, allUsers, outlines, preview, threshold);
                        if(exist != null){
                            mCallback.onRecognized(exist, transExtra);
                            exit();
                            hasDetectComplete = true;
                        }
                        mCallback.onFaceOutlineDetected(outlines);
                    }
                    if(timeout != 0){
                        if(!hasDetectComplete && SystemClock.uptimeMillis() - start >= timeout){
                            mCallback.onError("超时无人脸识别");
                            exit();
                        }
                    }
                }
            }catch (Exception e){
                exit();
            }
        }
        L.e("识别线程结束");
    }

    private List<AimoOutlineData> mapPoints(List<AimoRawFaceData> data){
        List<AimoOutlineData> ret = new ArrayList<>();
        for (AimoRawFaceData d : data){
            AimoOutlineData poi = new AimoOutlineData();
            poi.faceId = d.faceInfo.getId();
            poi.mapPoints(d.faceInfo.getPoints());
            poi.mapRect(d.faceInfo.getRect());
            ret.add(poi);
        }
        return ret;
    }

    public void exit(){
        isRunning = false;
        isExit.set(true);
    }

    public interface Callback{

        void onFaceOutlineDetected(List<AimoOutlineData> facePoints);

        void onRecognized(AimoFaceDbData user, Object... transExtra);

        void onError(String msg);
    }

    public static class Builder {

        private IRecognizeStrategy strategy;
        private Callback mCallback;
        private ImoImageFormat format;
        private ImoImageOrientation orientation;
        private Object[] transExtra;
        private float faceThreshold = 0.65f;
        private float timeoutSeconds = 0f;

        public Builder setStrategy(IRecognizeStrategy s) {
            strategy = s;
            return this;
        }

        public Builder callback(Callback callback) {
            mCallback = callback;
            return this;
        }

        public Builder format(ImoImageFormat format) {
            this.format = format;
            return this;
        }

        public Builder threshold(float threshold) {
            this.faceThreshold = threshold;
            return this;
        }

        public Builder orientation(ImoImageOrientation o) {
            this.orientation = o;
            return this;
        }

        public Builder timeout(float seconds){
            this.timeoutSeconds = seconds;
            return this;
        }

        public Builder transExtra(Object...extra){
            this.transExtra = extra;
            return this;
        }

        public RecognizeTask build() {
            return new RecognizeTask(this);
        }
    }
}
